from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import copy
import json

import os
from progress.bar import Bar
import torch
import matplotlib.pyplot as plt
import numpy as np
import random
from datetime import datetime
from torch._C import Graph
import time

from config import get_parse_args
from utils.logger import Logger
from utils.random_seed import set_seed
from datasets.dataset_factory import dataset_factory
from rl.model import load_model, save_model

from rl.env import Env
from rl.agent import Agent
from rl.config import RL_Config

# 导入可视化模块
from utils.visualization import draw_circuit, compare_circuits

selected_name_list = ['b04_C_000', 'b04_C_003', 'b04_C_005', 'b04_C_006', 
                      'b04_opt_C_002', 'b04_opt_C_006', 'b04_opt_C_007', 
                      'b07_C_000', 'b07_C_004', 'b07_C_011']

def train_RL(args):
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus_str

    print(args)
    args.num_rounds = args.test_num_rounds
    Logger(args)

    ####################################
    # 01 - Prepare data
    ####################################
    dataset = dataset_factory[args.dataset](args.data_dir, args)

    ####################################
    # 02 - Value Network
    ####################################
    config = RL_Config(args)
    agent = Agent(args, config)

    if args.resume:
        agent.Q_netT, _, start_epoch = load_model(
        agent.Q_netT, args.load_model, agent.optimizer, args.resume, args.lr, args.lr_step)
        agent.Q_netT.args.pretrain = False
        print(start_epoch)
        agent.Q_net.load_state_dict(agent.Q_netT.state_dict())

    if args.ftpt != 'no':
        if args.feature_pretrain_model == '':
            raise('No feature_pretrain_model filepath')
        agent.Q_netT.ft_net, _, ft_net_round = load_model(
        agent.Q_netT.ft_net, args.feature_pretrain_model, agent.optimizer, args.resume, args.lr, args.lr_step)
        agent.Q_netT.args.pretrain = False
        agent.Q_net.ft_net.load_state_dict(agent.Q_netT.ft_net.state_dict())
        print('[INFO] Read feature network with pre-trained round: ', ft_net_round)

    baseline_tc_list = []
    cop_tc_list = []
    dg_tc_list = []

    tot_begin_time = time.time()
    tmp_tp_each_round = args.no_tp_each_round

    # 创建可视化结果目录
    vis_dir = os.path.join(args.bench_dir, 'visualization')
    if not os.path.exists(vis_dir):
        os.makedirs(vis_dir)

    # 创建比较结果目录
    comp_dir = os.path.join(vis_dir, 'comparison')
    if not os.path.exists(comp_dir):
        os.makedirs(comp_dir)

    # 用于保存结果的字典
    visualization_results = {}

    for circuit_idx, g in enumerate(dataset):
        # if 'b17_C' not in g.name and 'mem_ctrl' not in g.name:
        #     continue
        # if 'b17' in g.name or 'mem_ctrl' in g.name:
        #     continue

        if 'b15' not in g.name:
            continue
        
        print(g.name)
        baseline_env = Env(g, config, args)
        dg_env = copy.deepcopy(baseline_env)
        cop_env = copy.deepcopy(baseline_env)

        # if baseline_env.baseline < 70:
        #     continue

        # Circuit information
        if args.circuit_info:
            print('[INFO] Circuit {}, # Nodes {:}, # PI {:}, # PO {:},  # Levels {:}'.format(
                g.name, len(g.forward_level), int(torch.sum(g.forward_level == 0)), int(torch.sum(g.backward_level == 0)), 
                int(max(g.forward_level))))
        
        # 记录当前电路信息
        circuit_results = {
            'name': g.name,
            'baseline_tc': 0,
            'cop_tc': 0,
            'dg_tc': 0,
            'baseline_bench': '',
            'cop_bench': '',
            'dg_bench': '',
            'visualization': {},
            'num_nodes': len(g.forward_level),
            'num_pi': int(torch.sum(g.forward_level == 0)),
            'num_po': int(torch.sum(g.backward_level == 0)),
            'levels': int(max(g.forward_level))
        }
        
        # Agent test
        if args.no_cp != -1:
            insert_cp_cnt = args.no_cp
        else:
            insert_cp_cnt = int(len(g.x) / 100)

        if (insert_cp_cnt / tmp_tp_each_round) > 20:
            args.no_tp_each_round = int(insert_cp_cnt / 20)
        else:
            args.no_tp_each_round = tmp_tp_each_round
        
        if tmp_tp_each_round == -1:
            args.no_tp_each_round = insert_cp_cnt

        print('\n===========================')
        print('Circuit {}'.format(g.name))
        print('Insert {:} Control Points'.format(insert_cp_cnt))
        print('Insert {:} TPs, {:} TPs / Step'.format(insert_cp_cnt, args.no_tp_each_round))
        agent.setInitState(dg_env.graph)
        terminal = 0
        dg_env.rename(dg_env.graph.name + '_dg')
        cp_idx = 1
        for cp_group in range(int(insert_cp_cnt / args.no_tp_each_round)):
            begin_time = time.time()
            cp_pos_list, cp_type_list = agent.getAction(dg_env.netlist, cp_idx, insert_cp_cnt)
            end_time = time.time()
            begin_cp_idx = cp_idx
            for idx in range(len(cp_pos_list)):
                cp_pos = cp_pos_list[idx]
                cp_type = cp_type_list[idx]
                dg_env.update_tp_list(cp_pos, cp_type)
                if not args.ignore_action:
                    print('[Action] Insert {:} CP on Pos {:}, Time: {:}s'.format(cp_type, cp_pos, end_time-begin_time))
                if len(cp_pos_list) + begin_cp_idx < insert_cp_cnt:
                    nextState, reward = dg_env.step_frame(cp_pos, cp_type, cp_idx, insert_cp_cnt)
                cp_idx += 1

        # COP
        cop_env.rename(cop_env.graph.name + '_cop')
        if not args.no_cop_tpi: 
            for cp_idx in range(insert_cp_cnt):
                cp_pos, cp_type = cop_env.netlist.cop_tpi(args, cp_idx)
                cop_env.update_tp_list(cp_pos, cp_type)
                nextState, reward = cop_env.step_frame(cp_pos, cp_type, cp_idx, insert_cp_cnt)

        # Print
        baseline_tc = baseline_env.get_test_coverage()
        dg_tc = dg_env.get_test_coverage()
        cop_tc = cop_env.get_test_coverage()
        print('Baseline: {:}%, COP: {:}%, DGRL: {:}%'.format(baseline_tc, cop_tc, dg_tc))

        # 更新电路结果
        circuit_results['baseline_tc'] = baseline_tc
        circuit_results['cop_tc'] = cop_tc
        circuit_results['dg_tc'] = dg_tc

        # if baseline_tc < 99:
        baseline_tc_list.append(baseline_tc)
        dg_tc_list.append(dg_tc)
        cop_tc_list.append(cop_tc)

        # Save
        if args.save_bench:
            # 保存bench文件
            baseline_bench_path = os.path.join(args.bench_dir, f"{g.name}.bench")
            cop_bench_path = os.path.join(args.bench_dir, f"{g.name}_cop.bench")
            dg_bench_path = os.path.join(args.bench_dir, f"{g.name}_dg.bench")
            
            baseline_env.netlist.save_bench(args, baseline_bench_path)
            cop_env.netlist.save_bench(args, cop_bench_path)
            dg_env.netlist.save_bench(args, dg_bench_path)

            # 记录bench文件路径
            circuit_results['baseline_bench'] = baseline_bench_path
            circuit_results['cop_bench'] = cop_bench_path
            circuit_results['dg_bench'] = dg_bench_path

            # 保存TP列表
            baseline_env.netlist.save_tp_list(args)
            cop_env.netlist.save_tp_list(args)
            dg_env.netlist.save_tp_list(args)
            
            baseline_env.netlist.save_verilog(args)
            cop_env.netlist.save_verilog(args)
            dg_env.netlist.save_verilog(args)
            print('[INFO] Save {}'.format(args.bench_dir + '/' + g.name + '.bench'))
            
            # 保存电路可视化结果
            try:
                # 基线电路可视化
                baseline_vis_path = os.path.join(vis_dir, f"{g.name}_baseline.png")
                draw_circuit(baseline_bench_path, baseline_vis_path)
                circuit_results['visualization']['baseline'] = baseline_vis_path
                
                # COP方法电路可视化
                cop_vis_path = os.path.join(vis_dir, f"{g.name}_cop.png")
                draw_circuit(cop_bench_path, cop_vis_path)
                circuit_results['visualization']['cop'] = cop_vis_path
                
                # DeepTPI方法电路可视化
                dg_vis_path = os.path.join(vis_dir, f"{g.name}_dg.png")
                draw_circuit(dg_bench_path, dg_vis_path)
                circuit_results['visualization']['dg'] = dg_vis_path
                
                # 创建对比图
                # 基线与COP方法比较
                baseline_cop_comp_path = os.path.join(comp_dir, f"{g.name}_baseline_vs_cop.png")
                compare_circuits(baseline_bench_path, cop_bench_path, 
                                baseline_cop_comp_path, 
                                title=f"Circuit {g.name}: Baseline vs COP (TC: {baseline_tc:.2f}% vs {cop_tc:.2f}%)")
                circuit_results['visualization']['baseline_vs_cop'] = baseline_cop_comp_path
                
                # 基线与DeepTPI方法比较
                baseline_dg_comp_path = os.path.join(comp_dir, f"{g.name}_baseline_vs_dg.png")
                compare_circuits(baseline_bench_path, dg_bench_path, 
                                baseline_dg_comp_path, 
                                title=f"Circuit {g.name}: Baseline vs DeepTPI (TC: {baseline_tc:.2f}% vs {dg_tc:.2f}%)")
                circuit_results['visualization']['baseline_vs_dg'] = baseline_dg_comp_path
                
                print(f"[INFO] Generated visualization images for {g.name}")
                
            except Exception as e:
                print(f"[WARNING] Failed to generate visualization for {g.name}: {str(e)}")
        
        # 保存电路结果
        visualization_results[g.name] = circuit_results

    tot_end_time = time.time()

    # 保存可视化结果概要
    results_summary = {
        'baseline_tc_avg': np.sum(baseline_tc_list)/len(baseline_tc_list) if baseline_tc_list else 0,
        'cop_tc_avg': np.sum(cop_tc_list)/len(cop_tc_list) if cop_tc_list else 0,
        'dg_tc_avg': np.sum(dg_tc_list)/len(dg_tc_list) if dg_tc_list else 0,
        'dg_improvement': np.average(np.array(dg_tc_list) - np.array(baseline_tc_list)) if dg_tc_list and baseline_tc_list else 0,
        'circuits': visualization_results,
        'runtime': tot_end_time-tot_begin_time
    }
    
    # 将结果保存为HTML报告
    html_path = os.path.join(vis_dir, "visualization_report.html")
    generate_html_report(results_summary, html_path)
    print(f"[INFO] Generated visualization report at {html_path}")

    # Average
    print('=======================================')
    print('Average Baseline TC: {:}%'.format(np.sum(baseline_tc_list)/len(baseline_tc_list)))
    print('Average COP TC: {:}%'.format(np.sum(cop_tc_list)/len(cop_tc_list)))
    print('Average DGRL TC: {:}%'.format(np.sum(dg_tc_list)/len(dg_tc_list)))

    print('*** Baseline ***')
    print(baseline_tc_list)
    print('*** COP ***')
    print(cop_tc_list)
    print('*** DGRL ***')
    print(dg_tc_list)

    print('Baseline, COP, DGRL')
    for idx in range(len(baseline_tc_list)):
        print('{:.4f}, {:.4f}, {:.4f}'.format(baseline_tc_list[idx]/100, cop_tc_list[idx]/100, dg_tc_list[idx]/100))

    print('*** DGRL Improvement ***')
    print(np.average(np.array(dg_tc_list) - np.array(baseline_tc_list)), '%')

    print('Total Run Time: {:}s'.format(tot_end_time-tot_begin_time))
    
    return results_summary

def generate_html_report(results, output_path):
    """
    生成HTML报告，展示测试结果和可视化图像
    
    参数:
        results: 测试结果数据
        output_path: HTML报告输出路径
    """
    html_content = f"""
    <!DOCTYPE html>
    <html>
    <head>
        <title>DeepTPI Test Results Visualization</title>
        <style>
            body {{ font-family: Arial, sans-serif; margin: 20px; }}
            h1, h2, h3 {{ color: #333366; }}
            .summary {{ background-color: #f0f0f0; padding: 15px; border-radius: 5px; margin-bottom: 20px; }}
            .circuit {{ border: 1px solid #ddd; padding: 15px; margin-bottom: 20px; border-radius: 5px; }}
            .circuit h3 {{ margin-top: 0; }}
            .metrics {{ display: flex; justify-content: space-around; margin-bottom: 15px; }}
            .metric {{ text-align: center; }}
            .metric .value {{ font-size: 24px; font-weight: bold; color: #333366; }}
            .metric .label {{ font-size: 14px; color: #666; }}
            .images {{ display: flex; flex-wrap: wrap; justify-content: space-around; }}
            .image-container {{ margin: 10px; text-align: center; }}
            .image-container img {{ max-width: 100%; max-height: 400px; border: 1px solid #ddd; }}
            .image-container .caption {{ margin-top: 5px; font-size: 14px; color: #666; }}
            table {{ border-collapse: collapse; width: 100%; margin-top: 20px; }}
            th, td {{ border: 1px solid #ddd; padding: 8px; text-align: left; }}
            th {{ background-color: #f2f2f2; }}
            tr:nth-child(even) {{ background-color: #f9f9f9; }}
        </style>
    </head>
    <body>
        <h1>DeepTPI Test Results Visualization</h1>
        
        <div class="summary">
            <h2>Summary</h2>
            <div class="metrics">
                <div class="metric">
                    <div class="value">{results['baseline_tc_avg']:.2f}%</div>
                    <div class="label">Average Baseline TC</div>
                </div>
                <div class="metric">
                    <div class="value">{results['cop_tc_avg']:.2f}%</div>
                    <div class="label">Average COP TC</div>
                </div>
                <div class="metric">
                    <div class="value">{results['dg_tc_avg']:.2f}%</div>
                    <div class="label">Average DeepTPI TC</div>
                </div>
                <div class="metric">
                    <div class="value">{results['dg_improvement']:.2f}%</div>
                    <div class="label">DeepTPI Improvement</div>
                </div>
                <div class="metric">
                    <div class="value">{results['runtime']:.2f}s</div>
                    <div class="label">Total Runtime</div>
                </div>
            </div>
            
            <h3>All Circuits</h3>
            <table>
                <tr>
                    <th>Circuit</th>
                    <th>Nodes</th>
                    <th>PI</th>
                    <th>PO</th>
                    <th>Levels</th>
                    <th>Baseline TC</th>
                    <th>COP TC</th>
                    <th>DeepTPI TC</th>
                    <th>Improvement</th>
                </tr>
    """
    
    # 添加每个电路的行
    for name, circuit in results['circuits'].items():
        improvement = circuit['dg_tc'] - circuit['baseline_tc']
        html_content += f"""
                <tr>
                    <td>{name}</td>
                    <td>{circuit['num_nodes']}</td>
                    <td>{circuit['num_pi']}</td>
                    <td>{circuit['num_po']}</td>
                    <td>{circuit['levels']}</td>
                    <td>{circuit['baseline_tc']:.2f}%</td>
                    <td>{circuit['cop_tc']:.2f}%</td>
                    <td>{circuit['dg_tc']:.2f}%</td>
                    <td>{improvement:.2f}%</td>
                </tr>
        """
    
    html_content += """
            </table>
        </div>
    """
    
    # 为每个电路添加详细信息
    for name, circuit in results['circuits'].items():
        html_content += f"""
        <div class="circuit">
            <h3>Circuit: {name}</h3>
            <div class="metrics">
                <div class="metric">
                    <div class="value">{circuit['baseline_tc']:.2f}%</div>
                    <div class="label">Baseline TC</div>
                </div>
                <div class="metric">
                    <div class="value">{circuit['cop_tc']:.2f}%</div>
                    <div class="label">COP TC</div>
                </div>
                <div class="metric">
                    <div class="value">{circuit['dg_tc']:.2f}%</div>
                    <div class="label">DeepTPI TC</div>
                </div>
                <div class="metric">
                    <div class="value">{circuit['dg_tc'] - circuit['baseline_tc']:.2f}%</div>
                    <div class="label">Improvement</div>
                </div>
            </div>
        """
        
        # 添加可视化图像
        if 'visualization' in circuit and circuit['visualization']:
            html_content += """
            <h4>Circuit Visualizations</h4>
            <div class="images">
            """
            
            # 添加基线、COP和DeepTPI方法的可视化
            visualization = circuit['visualization']
            if 'baseline' in visualization:
                rel_path = os.path.relpath(visualization['baseline'], os.path.dirname(output_path))
                html_content += f"""
                <div class="image-container">
                    <img src="{rel_path}" alt="Baseline Circuit">
                    <div class="caption">Baseline Circuit (TC: {circuit['baseline_tc']:.2f}%)</div>
                </div>
                """
            
            if 'cop' in visualization:
                rel_path = os.path.relpath(visualization['cop'], os.path.dirname(output_path))
                html_content += f"""
                <div class="image-container">
                    <img src="{rel_path}" alt="COP Circuit">
                    <div class="caption">COP Circuit (TC: {circuit['cop_tc']:.2f}%)</div>
                </div>
                """
            
            if 'dg' in visualization:
                rel_path = os.path.relpath(visualization['dg'], os.path.dirname(output_path))
                html_content += f"""
                <div class="image-container">
                    <img src="{rel_path}" alt="DeepTPI Circuit">
                    <div class="caption">DeepTPI Circuit (TC: {circuit['dg_tc']:.2f}%)</div>
                </div>
                """
            
            html_content += """
            </div>
            
            <h4>Circuit Comparisons</h4>
            <div class="images">
            """
            
            # 添加比较可视化
            if 'baseline_vs_cop' in visualization:
                rel_path = os.path.relpath(visualization['baseline_vs_cop'], os.path.dirname(output_path))
                html_content += f"""
                <div class="image-container">
                    <img src="{rel_path}" alt="Baseline vs COP">
                    <div class="caption">Baseline vs COP</div>
                </div>
                """
            
            if 'baseline_vs_dg' in visualization:
                rel_path = os.path.relpath(visualization['baseline_vs_dg'], os.path.dirname(output_path))
                html_content += f"""
                <div class="image-container">
                    <img src="{rel_path}" alt="Baseline vs DeepTPI">
                    <div class="caption">Baseline vs DeepTPI</div>
                </div>
                """
            
            html_content += """
            </div>
            """
        
        html_content += """
        </div>
        """
    
    html_content += """
    </body>
    </html>
    """
    
    # 保存HTML文件
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(html_content)

def test_specific_circuit(args, circuit_name):
    """
    测试一个特定的电路
    
    Args:
        args: 命令行参数
        circuit_name: 电路名称
        
    Returns:
        dict: 测试结果
    """
    try:
        print(f"[INFO] 测试电路: {circuit_name}")
        
        # 确保随机种子被设置
        if hasattr(args, 'random_seed'):
            random.seed(args.random_seed)
            torch.manual_seed(args.random_seed)
            np.random.seed(args.random_seed)
            
        # 初始化结果字典
        result = {"circuit": circuit_name}
        
        # 加载数据集
        print(f"[INFO] 数据目录: {args.data_dir}")
        
        # 添加安全加载类列表
        try:
            from torch_geometric.data.data import Data, DataEdgeAttr, DataTensorAttr
            from datasets.ordered_data import OrderedData
            from datasets.global_storage import GlobalStorage
            torch.serialization.add_safe_globals([OrderedData, GlobalStorage, DataEdgeAttr, DataTensorAttr])
            print("[INFO] 已添加安全加载类列表")
        except ImportError as e:
            print(f"[WARNING] 无法导入DataEdgeAttr或DataTensorAttr: {str(e)}")
            
        try:
            dataset = CircuitDataset(args.data_dir)
            print(f"[INFO] 成功加载数据集，包含 {len(dataset)} 个电路")
        except Exception as e:
            print(f"[ERROR] 加载数据集失败: {str(e)}")
            import traceback
            traceback.print_exc()
            return {"circuit": circuit_name, "error": f"加载数据集失败: {str(e)}"}
        
        # 初始化电路匹配变量
        found_exact_match = False
        found_approximate_match = False
        matching_circuit = None
        
        # 查找电路，包括精确和近似匹配
        for g in dataset:
            # 直接比较 - 精确匹配
            if g.name == circuit_name:
                matching_circuit = g
                found_exact_match = True
                print(f"[INFO] 找到精确匹配电路: {g.name}")
                break
                
            # 近似匹配 - 忽略大小写和特殊字符
            normalized_name = ''.join(c.lower() for c in g.name if c.isalnum())
            normalized_circuit_name = ''.join(c.lower() for c in circuit_name if c.isalnum())
            
            if normalized_name == normalized_circuit_name:
                matching_circuit = g
                found_approximate_match = True
                g.original_name = circuit_name  # 保存原始名称，用于结果报告
                print(f"[INFO] 找到近似匹配电路: {g.name} (与 {circuit_name} 匹配)")
                break
        
        # 如果没有找到精确匹配但找到了近似匹配
        if not found_exact_match and found_approximate_match:
            print(f"[WARNING] 找不到精确匹配电路: {circuit_name}, 使用近似匹配: {matching_circuit.name}")
        
        # 如果还是找不到，则尝试更宽松的匹配
        if matching_circuit is None:
            print(f"[WARNING] 找不到精确或近似匹配电路: {circuit_name}，尝试部分匹配")
            
            best_match = None
            best_score = 0
            
            for g in dataset:
                # 计算简单的部分匹配分数
                normalized_name = ''.join(c.lower() for c in g.name if c.isalnum())
                normalized_circuit_name = ''.join(c.lower() for c in circuit_name if c.isalnum())
                
                # 检查一个名称是否是另一个的子字符串
                if normalized_name in normalized_circuit_name or normalized_circuit_name in normalized_name:
                    score = len(set(normalized_name) & set(normalized_circuit_name))
                    if score > best_score:
                        best_score = score
                        best_match = g
            
            if best_match is not None:
                matching_circuit = best_match
                found_approximate_match = True
                matching_circuit.original_name = circuit_name  # 保存原始名称
                print(f"[INFO] 找到部分匹配电路: {matching_circuit.name} (与 {circuit_name} 部分匹配)")
            else:
                print(f"[ERROR] 无法找到与 {circuit_name} 匹配的电路")
                return {"circuit": circuit_name, "error": f"找不到匹配的电路: {circuit_name}"}
        
        # 使用匹配到的电路继续测试
        g = matching_circuit
                
        # 设置环境目录
        vis_dir = os.path.join(args.bench_dir, "visualizations")
        comp_dir = os.path.join(args.bench_dir, "comparisons")
        results_dir = os.path.join(args.bench_dir, "results")
        
        for dir_path in [vis_dir, comp_dir, results_dir]:
            if not os.path.exists(dir_path):
                os.makedirs(dir_path)
        
        # 加载预训练模型
        print(f"[INFO] 加载预训练模型: {args.feature_pretrain_model}")
        try:
            # 首先尝试使用weights_only=False加载
            try:
                checkpoint = torch.load(args.feature_pretrain_model, map_location=args.device, weights_only=False)
                model = TPIModel().to(args.device)
                model.load_state_dict(checkpoint)
                print("[INFO] 使用weights_only=False成功加载模型")
            except TypeError:
                # 如果失败，尝试使用默认参数加载
                checkpoint = torch.load(args.feature_pretrain_model, map_location=args.device)
                model = TPIModel().to(args.device)
                model.load_state_dict(checkpoint)
                print("[INFO] 使用默认参数成功加载模型")
        except Exception as e:
            print(f"[ERROR] 无法加载预训练模型: {str(e)}")
            import traceback
            traceback.print_exc()
            return {"circuit": circuit_name, "error": f"无法加载预训练模型: {str(e)}"}
        
        # 初始化环境
        baseline_env = ITC22Env(g, args, 0)
        dg_env = ITC22Env(g, args, 0)
        cop_env = ITC22Env(g, args, 0)
        
        # 初始化代理
        agent = DQNAgent(model, "test mode")
        
        # 确定插入的控制点数量
        if hasattr(g, 'num_gates'):
            insert_cp_cnt = 20
        else:
            insert_cp_cnt = 20
            print(f"[WARNING] 无法确定电路 {g.name} 的门数量，使用默认值 {insert_cp_cnt}")
        
        # 检查是否提供了--no_cp参数
        if hasattr(args, 'no_cp') and str(args.no_cp).strip() != '':
            try:
                insert_cp_cnt = int(args.no_cp)
                print(f"[INFO] 使用命令行参数设置的控制点数量: {insert_cp_cnt}")
            except ValueError:
                print(f"[WARNING] 无效的no_cp参数值: {args.no_cp}，使用默认值 {insert_cp_cnt}")
        
        # 确定每轮插入的TP数量
        tmp_tp_each_round = -1
        if hasattr(args, 'no_tp_each_round') and str(args.no_tp_each_round).strip() != '':
            try:
                tmp_tp_each_round = int(args.no_tp_each_round)
                print(f"[INFO] 使用命令行参数设置的每轮TP数量: {tmp_tp_each_round}")
            except ValueError:
                print(f"[WARNING] 无效的no_tp_each_round参数值: {args.no_tp_each_round}，使用默认值")
        
        if tmp_tp_each_round == 0:
            args.no_tp_each_round = int(insert_cp_cnt / 20)
        else:
            args.no_tp_each_round = tmp_tp_each_round
        
        if tmp_tp_each_round == -1:
            args.no_tp_each_round = insert_cp_cnt
        
        print('\n===========================')
        print('Circuit {}'.format(g.name))
        print('Insert {:} Control Points'.format(insert_cp_cnt))
        print('Insert {:} TPs, {:} TPs / Step'.format(insert_cp_cnt, args.no_tp_each_round))
        
        # DeepTPI测试
        agent.setInitState(dg_env.graph)
        dg_env.rename(dg_env.graph.name + '_dg')
        cp_idx = 1
        for cp_group in range(int(insert_cp_cnt / args.no_tp_each_round)):
            begin_time = time.time()
            cp_pos_list, cp_type_list = agent.getAction(dg_env.netlist, cp_idx, insert_cp_cnt)
            end_time = time.time()
            begin_cp_idx = cp_idx
            for idx in range(len(cp_pos_list)):
                cp_pos = cp_pos_list[idx]
                cp_type = cp_type_list[idx]
                dg_env.update_tp_list(cp_pos, cp_type)
                if not args.ignore_action:
                    print('[Action] Insert {:} CP on Pos {:}, Time: {:}s'.format(cp_type, cp_pos, end_time-begin_time))
                if len(cp_pos_list) + begin_cp_idx < insert_cp_cnt:
                    nextState, reward = dg_env.step_frame(cp_pos, cp_type, cp_idx, insert_cp_cnt)
                cp_idx += 1
        
        # COP测试
        cop_env.rename(cop_env.graph.name + '_cop')
        if not args.no_cop_tpi: 
            for cp_idx in range(insert_cp_cnt):
                cp_pos, cp_type = cop_env.netlist.cop_tpi(args, cp_idx)
                cop_env.update_tp_list(cp_pos, cp_type)
                nextState, reward = cop_env.step_frame(cp_pos, cp_type, cp_idx, insert_cp_cnt)
        
        # 获取测试覆盖率
        baseline_tc = baseline_env.get_test_coverage()
        dg_tc = dg_env.get_test_coverage()
        cop_tc = cop_env.get_test_coverage()
        improvement = dg_tc - baseline_tc
        
        print('Baseline: {:}%, COP: {:}%, DGRL: {:}%'.format(baseline_tc, cop_tc, dg_tc))
        print('Improvement: {:}%'.format(improvement))
        
        # 更新结果
        result['baseline_tc'] = baseline_tc
        result['cop_tc'] = cop_tc
        result['dg_tc'] = dg_tc
        result['improvement'] = improvement
        
        # 如果使用了近似匹配的电路，添加原始电路名称到结果
        if hasattr(g, 'original_name'):
            result['original_name'] = g.original_name
        
        # 保存结果到文件
        result_file = os.path.join(results_dir, f"{circuit_name}_results.json")
        with open(result_file, 'w', encoding='utf-8') as f:
            json.dump({
                "circuit": circuit_name,
                "baseline_tc": f"{baseline_tc:.2f}%",
                "cop_tc": f"{cop_tc:.2f}%",
                "dg_tc": f"{dg_tc:.2f}%",
                "improvement": f"{improvement:.2f}%",
                "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
            }, f, indent=2, ensure_ascii=False)
        
        # 保存bench文件
        if args.save_bench:
            baseline_bench_path = os.path.join(args.bench_dir, f"{g.name}.bench")
            cop_bench_path = os.path.join(args.bench_dir, f"{g.name}_cop.bench")
            dg_bench_path = os.path.join(args.bench_dir, f"{g.name}_dg.bench")
            
            baseline_env.netlist.save_bench(args, baseline_bench_path)
            cop_env.netlist.save_bench(args, cop_bench_path)
            dg_env.netlist.save_bench(args, dg_bench_path)
            
            # 保存TP列表
            baseline_env.netlist.save_tp_list(args)
            cop_env.netlist.save_tp_list(args)
            dg_env.netlist.save_tp_list(args)
            
            print('[INFO] Save {}'.format(args.bench_dir + '/' + g.name + '.bench'))
            
            # 生成可视化
            try:
                # 基线电路可视化
                baseline_vis_path = os.path.join(vis_dir, f"{g.name}_baseline.png")
                draw_circuit(baseline_bench_path, baseline_vis_path)
                
                # COP方法电路可视化
                cop_vis_path = os.path.join(vis_dir, f"{g.name}_cop.png")
                draw_circuit(cop_bench_path, cop_vis_path)
                
                # DeepTPI方法电路可视化
                dg_vis_path = os.path.join(vis_dir, f"{g.name}_dg.png")
                draw_circuit(dg_bench_path, dg_vis_path)
                
                # 创建比较图
                comp_path = os.path.join(comp_dir, f"{g.name}_comparison.png")
                compare_circuits(baseline_bench_path, cop_bench_path, dg_bench_path, comp_path)
                
                print(f"[INFO] 已生成电路可视化图片")
            except Exception as e:
                print(f"[WARNING] 生成可视化图片失败: {str(e)}")
        
        return result
    except Exception as e:
        print(f"[ERROR] 测试电路 {circuit_name} 时出错: {str(e)}")
        import traceback
        traceback.print_exc()
        return {"circuit": circuit_name, "error": str(e)}

if __name__ == '__main__':
    args = get_parse_args()
    set_seed(args)
    random.seed(int(datetime.now().timestamp()))
    
    # 检查RL_model参数
    if not hasattr(args, 'RL_model') or args.RL_model is None:
        print("[INFO] 未找到RL_model参数，设置为默认值'non_level'")
        args.RL_model = 'non_level'
    
    # 检查CUDA是否可用，如果不可用，则强制使用CPU
    if torch.cuda.is_available() and hasattr(args, 'gpus') and args.gpus >= 0:
        args.device = torch.device(f'cuda:{args.gpus}')
    else:
        args.device = torch.device('cpu')
        print("[INFO] CUDA不可用或未启用，使用CPU运行")
    
    # 如果指定了单个电路，则只测试该电路
    if hasattr(args, 'circuit') and args.circuit and args.circuit != "":
        print(f"[INFO] 测试单个电路: {args.circuit}")
        result = test_specific_circuit(args, args.circuit)
        print(f"测试结果: {result}")
    else:
        # 否则测试所有电路
        print("[INFO] 测试所有电路")
        results = train_RL(args)

